# YOLOv5 🚀 by Ultralytics, GPL-3.0 license
"""
Common modules
"""

import json
import math
import platform
import warnings
from collections import OrderedDict, namedtuple
from copy import copy
from pathlib import Path

import cv2
import numpy as np
import pandas as pd
import requests
import torch
import torch.nn as nn
# from PIL import Image
# from torch.cuda import amp
#
# from utils.datasets import exif_transpose, letterbox
# from utils.general import (LOGGER, check_requirements, check_suffix, check_version, colorstr, increment_path,
#                            make_divisible, non_max_suppression, scale_coords, xywh2xyxy, xyxy2xywh)
# from utils.plots import Annotator, colors, save_one_box
# from utils.torch_utils import copy_attr, time_sync
from models.common import Conv


class C2fBottleneckPruned(nn.Module):
    # Standard bottleneck
    def __init__(self, cv1in, cv1out, cv2out, shortcut=True, g=1, k=(3, 3), e=0.5):  # ch_in, ch_out, shortcut, groups, kernels, expand
        super().__init__()
        c_ = int(cv2out * e)  # hidden channels
        self.cv1 = Conv(cv1in, cv1out, k[0], 1)
        self.cv2 = Conv(cv1out, cv2out, k[1], 1, g=g)
        self.add = shortcut and cv1in == cv2out

    def forward(self, x):
        return x + self.cv2(self.cv1(x)) if self.add else self.cv2(self.cv1(x))


class BottleneckPruned(nn.Module):
    # Pruned bottleneck
    def __init__(self, cv1in, cv1out, cv2out, shortcut=True, g=1):  # ch_in, ch_out, shortcut, groups, expansion
        super(BottleneckPruned, self).__init__()
        self.cv1 = Conv(cv1in, cv1out, 1, 1)
        self.cv2 = Conv(cv1out, cv2out, 3, 1, g=g)
        self.add = shortcut and cv1in == cv2out

    def forward(self, x):
        return x + self.cv2(self.cv1(x)) if self.add else self.cv2(self.cv1(x))


class c2fPruned(nn.Module):
    # CSP Bottleneck with 2 convolutions
    def __init__(self, cv1in, cv1out, cv2out, bottle_args, n=1, shortcut=False, g=1, e=0.5):  # ch_in, ch_out, number, shortcut, groups, expansion
        super().__init__()
        # self.c = int(cv2out * e)  # hidden channels
        cv2in = bottle_args[-1][-1]
        self.cv1 = Conv(cv1in, cv1out, 1, 1)
        self.cv2 = Conv(cv2in+cv1out, cv2out, 1)  # optional act=FReLU(c2)
        self.m = nn.Sequential(*[C2fBottleneckPruned(*bottle_args[i], shortcut, g, k=((3, 3), (3, 3)), e=1.0) for i in range(n)])

    def forward(self, x):
        # torch.chunk在给定维度上对张量进行分块 -> torch.chunk(2, 1) 在第一维（通道）上将张量拆成俩份
        y = list(self.cv1(x).chunk(2, 1))
        # list.extend 在列表list的末尾添加元素
        # y.extend(m(y[-1])for m in self.m) # 依次将分割特征图cv1后的子特征图送入模型的各个模块中
        for m in self.m:
            # print(f"m:{m}")
            print(f" ")
            a = m(y[-1])
            y.extend(a)
            # print(a)
        return self.cv2(torch.cat(y, 1)) # 在第一个维度（通道）进行相加

    # def forward_split(self, x):
    #     y = list(self.cv1(x).split((self.c, self.c), 1))
    #     y.extend(m(y[-1]) for m in self.m)
    #     return self.cv2(torch.cat(y, 1))


class C3Pruned(nn.Module):
    # CSP Bottleneck with 3 convolutions
    def __init__(self, cv1in, cv1out, cv2out, cv3out, bottle_args, n=1, shortcut=True, g=1):  # ch_in, ch_out, number, shortcut, groups, expansion
        super(C3Pruned, self).__init__()
        cv3in = bottle_args[-1][-1]
        self.cv1 = Conv(cv1in, cv1out, 1, 1)
        self.cv2 = Conv(cv1in, cv2out, 1, 1)
        self.cv3 = Conv(cv3in+cv2out, cv3out, 1)
        # bottle_args = [[16, 16, 16]]
        self.m = nn.Sequential(*[BottleneckPruned(*bottle_args[k], shortcut, g) for k in range(n)])

    def forward(self, x):
        return self.cv3(torch.cat((self.m(self.cv1(x)), self.cv2(x)), dim=1))


class SPPFPruned(nn.Module):
    # Spatial pyramid pooling layer used in YOLOv3-SPP
    def __init__(self, cv1in, cv1out, cv2out, k=5):
        super(SPPFPruned, self).__init__()
        self.cv1 = Conv(cv1in, cv1out, 1, 1)
        self.cv2 = Conv(cv1out * 4, cv2out, 1, 1)
        self.m = nn.MaxPool2d(kernel_size=k, stride=1, padding=k // 2)

    def forward(self, x):
        x = self.cv1(x)
        with warnings.catch_warnings():
            warnings.simplefilter('ignore')  # suppress torch 1.9.0 max_pool2d() warning
            y1 = self.m(x)
            y2 = self.m(y1)
            return self.cv2(torch.cat([x, y1, y2, self.m(y2)], 1))
